# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import functools

from coin.base.datetime_util import to_datetime
from coin.flow.topic_map_util import get_worker_id


def _other_worker(worker_id):
  if worker_id == '1':
    return '2'
  elif worker_id == '2':
    return '1'
  else:
    raise ValueError('Invalide worker_id: %s' % worker_id)


class DoubleShiftHandler(object):
  def __init__(self, kls):
    self._handlers = {'1': kls(), '2': kls()}
    self._timestamps = {'1': 0, '2': 2}
    self._master = None

  def subscribe(self, symbol, callback):
    for key, value in self._handlers.items():
      value.subscribe(symbol, callback)

  def _check_timeout(self, worker_id):
    if self._master and worker_id != self._master:
      diff = self._timestamps[worker_id] - self._timestamps[self._master]
      if diff >= 5 * 1000 * 1000 * 1000:  # Nanoseonds
        self._master = worker_id

  def get_handler(self, protocol: str):
    # TODO(inkyu): Implement BaseHandler
    return functools.partial(self.update_book, protocol)

  def update_book(
      self,
      protocol: str,  # TODO(inkyu): Make me enum
      record,
      queue_data,
      topic_data,
      subscribe=True):
    timestamp = record.timestamp
    if protocol == 'rest':
      # Update both handlers for rest message
      for worker_id in self._handlers:
        handler = self._handlers[worker_id]
        books = handler.update_book(protocol, record, queue_data, topic_data)
        if handler.ready:
          handler.handle(books)
    elif protocol == 'ws':
      worker_id = get_worker_id(topic_data.topic_id)
      assert self._timestamps[worker_id] < timestamp
      self._timestamps[worker_id] = timestamp
      books = self._handlers[worker_id].update_book(protocol, record, queue_data, topic_data)
      handler = self._handlers[worker_id]
      if not handler.ready:
        # Dead
        other_handler = self._handlers[_other_worker(worker_id)]
        if other_handler.ready:
          self._master = _other_worker(worker_id)
        else:
          # Both books are dead
          pass
      else:
        if not self._master and self._handlers[worker_id].ready:
          self._master = worker_id
        self._check_timeout(worker_id)
        # Alive & master
        if self._master == worker_id:
          handler.handle(books)
    else:
      raise ValueError('Invalid protocol: %s' % protocol)
